import paddle
import os
from paddle.inference import Config,create_predictor,PrecisionType

#设备选取
def get_infer_gpuid():
    if os.name == 'nt':
        try:
            return int(os.environ['CUDA_VISIBLE_DEVICES'].split(',')[0])
        except KeyError:
            return 0
    if not paddle.fluid.core.is_compiled_with_rocm():
        cmd = "env | grep CUDA_VISIBLE_DEVICES"
    else:
        cmd = "env | grep HIP_VISIBLE_DEVICES"
    env_cuda = os.popen(cmd).readlines()
    if len(env_cuda) == 0:
        return 0
    else:
        gpu_id = env_cuda[0].strip().split("=")[1]
        return int(gpu_id[0])


#构建推理引擎
def Init_Paddle(model_dir,infer_type="fp16",use_tensorrt=True):
    '''
    :parma model_dir: 模型路径
    :param infer_type: 推理以什么类型进行,如fp32、fp16、int8等
    :param use_tensorrt: 是否使用tensorrt进行推理
    '''

    gpu_id = get_infer_gpuid()
    if gpu_id is None:
        print("GPU is not found in current device by nvidia-smi. Please check your device or ignore it if run on jeston.")

    #模型文件
    model_path=os.path.join(model_dir,"model.pdmodel")
    #权重文件
    weights_path=os.path.join(model_dir,"model.pdiparams")
    config=Config(model_path,weights_path)


    #数据类型确定
    if infer_type=="fp32":
        precision=PrecisionType.Float32
    elif infer_type=="fp16":
        precision=PrecisionType.Half
    else:
        precision=PrecisionType.Int8

    # 设置使用gpu,参数一是初始化分配的gpu显存，以MB为单位；参数二是设备名称
    config.enable_use_gpu(500, 0)
    # 使用tensorrt
    if use_tensorrt:
        '''
        workspace_size     - 指定 TensorRT 使用的工作空间大小
        max_batch_size     - 设置最大的 batch 大小，运行时 batch 大小不得超过此限定值
        min_subgraph_size  - Paddle-TRT 是以子图的形式运行，为了避免性能损失，当子图内部节点个数
                            大于 min_subgraph_size 的时候，才会使用 Paddle-TRT 运行
        precision          - 指定使用 TRT 的精度，支持 FP32(kFloat32)，FP16(kHalf)，Int8(kInt8)
        use_static         - 若指定为 true，在初次运行程序的时候会将 TRT 的优化信息进行序列化到磁盘上，
                            下次运行时直接加载优化的序列化信息而不需要重新生成
        use_calib_mode     - 若要运行 Paddle-TRT INT8 离线量化校准，需要将此选项设置为 true
        '''
        config.enable_tensorrt_engine(
            workspace_size=1 << 30,
            precision_mode=precision,
            max_batch_size=1,
            use_static=True,
            min_subgraph_size=15)

    #enable_memory_optim推理优化
    config.enable_memory_optim()
    #预测时不出现日志
    config.disable_glog_info()

    config.switch_use_feed_fetch_ops(False)
    config.switch_ir_optim(True) #是否启用IR优化

    #创建预测器
    predictor=create_predictor(config)

    #输入
    input_names=predictor.get_input_names()
    for name in input_names:
        input_tensor=predictor.get_input_handle(name)

    #输出
    output_names=predictor.get_output_names()
    output_tensors=[]
    for output_name in output_names:
        output_tensor=predictor.get_output_handle(output_name)
        output_tensors.append(output_tensor)

    return predictor, input_tensor, output_tensors, config

